package com.example.demo;

import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@RestController
public class EmbeddingController {

    private final EmbeddingModel embeddingModel;

    @Autowired
    VectorStore vectorStore;

    @Autowired
    public EmbeddingController(EmbeddingModel embeddingModel) {
        this.embeddingModel = embeddingModel;
    }

    @GetMapping("/ai/embedding")
    public String embed(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        EmbeddingResponse embeddingResponse = this.embeddingModel.embedForResponse(List.of(message));
        return Arrays.toString(embeddingResponse.getResult().getOutput());
    }

    @GetMapping("/ai/testEmbedding")
    public String queryEmbedding(){
        try {
            List <Document> documents = List.of(
                    new Document("重庆大学虎溪校区竹五经常停水，让小伙伴们苦不堪言", Map.of("meta1", "meta1")));

            vectorStore.add(documents);

            SearchRequest searchRequest = SearchRequest.builder().query("重庆大学").topK(5).build();
            List<Document> results = this.vectorStore.similaritySearch(searchRequest);
            if (results != null) {
                StringBuilder stringBuilder = new StringBuilder();
                for (Document result : results) {
                    stringBuilder.append(result.getText()).append("\n");
                }
                return stringBuilder.toString();
            }else{
                return "空空如也";
            }
        } catch (Exception e) {
            return e.getMessage();
        }
    }

    @PostMapping("/add")
    public String add() {

        List<Document> documents = List.of(
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
                new Document("The World is Big and Salvation Lurks Around the Corner"),
                new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));

        vectorStore.add(documents);
        List<Document> results = vectorStore.similaritySearch("Spring");

        return results.stream().map(Document::getText).collect(Collectors.joining("\n"));
    }


}
